sequenceDiagram
  participant SchedulerThread
  participant CompletionFlags
  participant WakeList
  participant Watermark

  SchedulerThread->>CompletionFlags: try_set_completion_flag(thread_idx, local_id)
  alt flag set successfully
    CompletionFlags-->>SchedulerThread: success
    SchedulerThread->>WakeList: drain_wake_list(thread_idx)
    SchedulerThread->>Watermark: update_completed_watermark(thread_idx, local_id)
  else reuse not yet certified
    CompletionFlags-->>SchedulerThread: failure
    SchedulerThread->>SchedulerThread: push to failed_heap_of_set_completion_flag
    SchedulerThread->>SchedulerThread: retry_set_completion_flags(thread_idx) later
  end

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

_{Comment @coderabbitai help to get the list of available commands.}

Thanks for putting both variants up side by side — having #1618 and #1619 as an explicit either/or made this much easier to reason about.

Two separate topics below: the forward-looking functional groundwork, and the performance case.

1. Functional groundwork (flag-slot reuse) — the premise does not hold in hbg

The int32 completion stamps, the slot-reuse gate, the failsafe min-heap, and cached_completed_watermark all exist to support flag-slot reuse. Reuse only becomes useful once the ring actually wraps, and that cannot happen in host_build_graph today — not because reclaim is unimplemented, but because of the execution model:

hbg's orchestrator and scheduler are strictly serial, not overlapped.

host:   run_host_orchestration()   <- the entire graph is built here
host:   copy_to_device()           <- the populated SM is uploaded
--------------------------------------------------------------------
device: AICPU kernel boots scheduler-only
        ("host_build_graph has no device-side orchestrator", aicpu_executor.cpp)

By the time any scheduler thread exists, the host has already finished submitting. So even with a complete reclaim protocol in place, the slots freed by reclaim would have no consumer — nobody is waiting to submit into them. Reclaim pays off only when the host wants to submit task N while the window is full, i.e. when O and S overlap in time. That is the tensormap_and_ringbuffer model, not hbg's.

Consistent with that, last_task_alive never advances here:

// pto_scheduler.h:26
// reclaim (whole-graph-resident), so last_task_alive is not advanced here.

which pins the ring allocator's gate at local_task_id_ + 1 < window_size_. We hit that live while benchmarking — paged_attention Case1 (~65792 tasks) on the default 16384 window:

[TaskAllocator] BLOCKED: tasks=16383/16384, on=task, spins=350720000
Assertion failed: index < output_count_   (pto_types.h:114)

(It runs fine with PTO2_RING_TASK_WINDOW=131072 PTO2_RING_HEAP=805306368.)

One more thing worth noting for whenever reclaim is taken up: completion_flags is one of four parallel arrays indexed by local_id & task_window_mask — alongside task_descriptors, task_payloads, and slot_states. Wrapping aliases all four. The other three hold live wake lists, fanin arrays other tasks are still reading, and in-flight descriptors, so they need a real reclaim protocol rather than an id stamp. Landing the flags piece on its own does not move the task-count ceiling.

So as things stand, none of this machinery is exercised on hbg: the reuse gate never fails, the failsafe heap stays empty, and cached_completed_watermark only mitigates a watermark fallback that the reuse itself introduced.

2. Performance — we measure this PR at parity with its base

We benchmarked the PR head against its own merge-base on one card.

Results (a2a3, aicpu_thread_num=4, device wall-clock, 100 rounds/side)

bgemm / matmul / vector_example were also run, but at 0.10–0.19 ms their measured cv is 5–12%, so they carry no resolution at this scale and we are not quoting them.

Method

• Baseline = this PR's merge-base 80aa2876; current = a3a637f0. The only difference between the two trees is this PR's single commit.
• Two detached-HEAD worktrees, each with its own venv and pip install --no-build-isolation -e ., so the nanobind extension and the runtime binaries are fully isolated per side.
• Same physical card: the whole comparison runs inside a single task-submit invocation, so one device lock covers both sides; baseline and PR run back-to-back per workload.
• 100 rounds per side, round 0 dropped as warm-up, outliers trimmed, then the mean.
• Metric is the runner_run.device_wall span from the [STRACE] markers, rendered by python -m simpler_setup.tools.strace_timing <log> --rounds-table.
• On pa, round times fall into two groups (~24 ms and ~51 ms, roughly half each, on both sides); the figure quoted above is the fast group. We are looking into the cause separately.
• paged_attention Case1 needs PTO2_RING_TASK_WINDOW=131072 PTO2_RING_HEAP=805306368 to run at all (see §1).
• qwen3_14b_decode ships as runtime="tensormap_and_ringbuffer"; we re-pointed it to "host_build_graph" — a one-line change, verified byte-identical in both worktrees.

Cross-check against #1544

Our base (which contains 3S+1P) measures 24.816 ms on pa Case1. The 3S+1P number posted in #1544 is 25.013 ms — 0.8% apart, from an independent measurement on a different device on a different day.

Taking #1544's 4S baseline (28.016 ms) as the reference for a plain revert:

plain 4S revert (no optimisations)   28.016 ms   [#1544]
#1618 (4S + shuffle + frontier gate) 25.328 ms   [this measurement]
3S+1P base                           24.816 ms   [this measurement]

So the bit-shuffle and the frontier-gated watermark clearly do real work — they recover roughly 9.6 of the 10.7 points a plain revert would cost (that comparison is cross-experiment — different day, device, and 4S baseline build — so treat the figure as indicative). But against the current base they land at parity, not ahead.

What we would like to reconcile

Our numbers do not line up with the improvement the charts in the PR body appear to show, and we cannot read the underlying values off the images. Could you post the numbers as text, along with your measurement method and what the baseline is — main, or #1619? Since both PRs revert 3S+1P, a #1618-vs-#1619 delta measures only eager-vs-lazy watermark maintenance, which is a much narrower claim than #1618-vs-main.

3. 3S+1P is not only about ready-queue contention

One consideration behind the 3S+1P split that we would like to see represented in the comparison: the dedicated P thread does more than remove multi-producer contention. It also takes async deferred-completion polling and dummy / predicate-failed retirement off the dispatch path entirely. From main's own comment, which this PR removes:

Async deferred-completion polling and dependency-only (dummy / predicate-failed) retirement both run on P, which owns every completion→ready transition — the scheduler threads' loop stays purely core-local (poll own COND, dispatch own cores) and never touches the shared mailbox or dummy queue.

and, for the dummy queue specifically:

P produces and drains it, so the queue is single-threaded end to end.

Under 4S both come back onto every scheduler thread's loop: async_wait_list.poll_and_complete behind a shared try_lock(), and dummy_ready_queue becomes a contended MPMC queue drained by all four threads. For graphs with substantial async or dummy/barrier traffic, that is work now sitting on the dispatch path of every thread rather than on a thread that owns no cores. We expect 3S+1P to hold a substantial advantage on those workloads that 4S cannot match, and that class is where the two designs should really be separated.

To be clear on where we stand: §1 and §3 are observations on the design, and §2 is a request for data rather than an objection. If the numbers show #1618 ahead of main under a method we can reproduce, that settles it.

1. Functional groundwork (flag-slot reuse) — the premise does not hold in hbg

There are two options here:

• The number of tasks is always less than (or equal to) the number of flags.
• The number of tasks may be larger the number of flags.

In the first case, there is no need for the mask in local_id & task_window_mask. Furthermore, I would also question the need for the watermark in the first place? As this would remove the cache-line contention.

In the second case, reuse of flags is necessary. This also removes the necessity to increase PTO2_RING_TASK_WINDOW (at least from the perspective of the completion flags).

If the first case is truely the design choice, then a lot of the PR can be stripped.

3. 3S+1P is not only about ready-queue contention

I believe it all comes down to a proper benchmark for this case.
In PyPTO 2, this is something we investigated and we have found the following:

• The benefit of the single threaded queue does not outweigh the latency added by moving task responsibilities through SPSC queues to another thread.
• Having 4 instead of 3 scheduling threads drastically improves reaction time to changes in the AICores and increases their uptime.
• Due to batching, the MPMC queues are not that expensive.

I suggest the different designs 3+1, 4, and 4+1, should all be investigated on a proper benchmark.

Thanks — that is exactly the right way to frame it, and the answer is case 1.

Case 1 is not a preference, it is enforced today

The ring allocator gate is local_task_id_ - last_alive + 1 < window_size_, and last_task_alive never advances in hbg (pto_scheduler.h:26), so it reduces to local_task_id_ + 1 < window_size_. Task count is therefore hard-capped at window - 1, while completion_flags is window entries long. tasks <= flags holds by construction — case 2 is not something the runtime can currently reach, and hitting the cap is a hard failure, not a wrap:

[TaskAllocator] BLOCKED: tasks=16383/16384, on=task, spins=350720000

Why the code reads ambiguously here

host_build_graph began as a port of tensormap_and_ringbuffer — the tensormap, the ring buffer, and the orchestrator+scheduler split were all carried over together, deliberately, to keep the two runtimes structurally close. But hbg's defining constraint is different: the orchestrator must run to completion before the scheduler starts, and that is intentional, not a limitation we plan to lift.

The result is that several ring concepts survive in hbg without carrying any meaning:

• last_task_alive is never advanced, so the ring never wraps;
• task_window_mask never actually masks anything, because local_id cannot exceed the window;
• the reclaim machinery exists but has nothing to reclaim on behalf of.

These vestiges are what make the design intent ambiguous when reading the code, and removing them is precisely what we want to do next. So your instinct that a lot can be stripped is right — and it reaches beyond the code this PR touches.

It is also the direction, not just the status quo

The concern behind case 2 is presumably that a large model submits too many tasks — a 40-layer Qwen decode being the obvious example. Our answer to that is not a bigger flag array but a smaller task count: #1444 adds a GRAPH composite task type that records a repeated sub-DAG once and then submits one outer ring task per invocation, with the internal nodes held in AICPU-local descriptor/payload/slot storage that consumes no ring slot at all. A 40-layer decode becomes 40 ring tasks rather than 40×N.

So the intent is to keep task counts comfortably below the flag array, not to make the array circular.

On the two things you would strip

The mask — agreed, under case 1 local_id & task_window_mask is redundant.

The watermark — you are right to question it, and it holds up on inspection. Once reuse is out, its only reader is wait_one_consumers inside wait_for_tensor_ready(..., wait_for_consumers=true), and that path is effectively unreachable in hbg. flush_segment runs wait_one_producer first and returns on failure, so the watermark is only read once the producer's task_state has already reached COMPLETED. A device task cannot reach that during orchestration — O runs to completion before S starts — so those callers time out at the producer wait and never get there. The only producers that do pass are the ones the host completed inline, and for those the watermark check is either trivially satisfied or a timeout in its own right.

So it performs no real gating today. The one thing worth deciding deliberately rather than as a side effect: removing it also removes set_tensor_data's WAR gate, which is public API surface even in its currently non-functional state.

One correction on case 2

This also removes the necessity to increase PTO2_RING_TASK_WINDOW (at least from the perspective of the completion flags).

The parenthetical is doing a lot of work there. task_window_size sizes four parallel arrays — task_descriptors, task_payloads, slot_states, completion_flags — and it also sets the allocator gate above. Flag reuse on its own therefore would not remove the need to size the window to the task count; the other three arrays and the gate still require it. With GRAPH tasks in place, the task count stays low enough that the default window is sufficient anyway.

Agreed — a proper benchmark is the right next step, and it is what we intend to focus on next. A shared benchmark is also what we need for performance alignment generally, not just for this decision.

Where the thread budget comes from

Framing this as "3 vs 4 scheduling threads" understates what actually changed.

The AICPU thread count is an architecture-level choice made per platform, for hardware reasons we are happy to go into separately:

• a2a3: 4 threads, spent in tensormap_and_ringbuffer as 1 orchestrator + 3 schedulers — the orchestrator runs on the AICPU and owns no cores;
• a5: 5 threads, as 1O + 4S.

host_build_graph moves the orchestrator to the host, which frees exactly one thread. On a2a3 we spent that freed thread on a dedicated resolution thread, giving 3S + 1P.

So the number of core-owning threads has been 3 throughout on a2a3. What #1618 does is not "3 → 4 schedulers"; it is spending the thread freed by moving O to the host on a fourth core-owner rather than on a resolver. The real question is therefore how best to spend that freed thread, and that is exactly what a benchmark should answer.

On 4+1

4+1 needs five scheduling threads, so it is not measurable on a2a3, where the budget is four — there the comparison is 3S+1P vs 4S. a5's budget is five, but how it gets spent is a separate discussion from this PR.

On the MPMC point

That is plausible and we do not dispute it — batching amortises the enqueue cost, and our own dispatch path already pops in batches sized to the thread's free cores. If MPMC contention is not the dominant term, then the case for a dedicated resolver rests on the other half of what P does, which is where we would want the benchmark to look.

Where the data currently stands

From our measurements: on ordinary workloads, 3S+1P and #1618's 4S are at parity — qwen and paged_attention both land within noise. So on the cases we can measure today, the two designs are equivalent, and the decision cannot be made on those numbers alone.

Our expectation — and it is an expectation, not a measurement — is that the separation shows up on async-heavy graphs, where P absorbs the mailbox polling and dependency-only retirement that otherwise sit on every scheduler thread's dispatch loop. No current hbg test exercises async at all, so that regime is simply unmeasured today.

Next step

We will drive building the benchmark set. To be useful it needs to span both ends — scheduler-throughput-bound and compute-bound — and add async/dummy-heavy cases, which is the gap.

@noabauma — flagging you since you own the performance testing.

Thanks for the charts — having the ladder made this straightforward to reproduce. We ran the same ladder against this PR's merge-base and cannot reproduce the improvement: we measure the two branches at parity on every rung.

Our results — device_wall, base (80aa287) vs #1618 (a3a637f)

Every rung is within ±0.3%, against cv of 1.0–9.5%.

Method

• base = this PR's merge-base 80aa2876, pr = a3a637f0; the only difference between the two trees is this PR's single commit.
• Two worktrees with isolated venvs. We verified the two sides really are different builds: the libaicpu_kernel.so binaries differ, base contains run_resolution_thread and no cached_completed_watermark, pr the reverse.
• Cases from Add: host_build_graph copies of the paged_attention examples #1667's examples/a2a3/host_build_graph/paged_attention, unmodified — Case1 and Case2 keep its runtime_env verbatim (2 GB and 1 GB heap). Mid4 / Mid16 / Mid64 reconstructed as batch 4 / 16 / 64 with Case1's other parameters; the resulting task counts (1,028 / 4,112 / 16,448) match your ladder. bgemm from tests/st/a2a3/host_build_graph/bgemm.
• One device, one lock: all twelve runs execute inside a single task-submit invocation on device 8, so both branches share the same card for the whole comparison, and base / pr run back-to-back within each workload.
• Rounds matched to your reported n. Round 0 dropped as warm-up. Metric is the runner_run.device_wall span from the [STRACE] markers.

Side by side

Our base tracks your main closely on Mid4 / Mid16 / Mid64 (within 1–6%), which suggests the cases, the ring configuration and the platform are aligned. The divergence is concentrated on the two largest rungs.

Where your figures fall inside our distributions:

What we think could produce this

1. The runs may not have been interleaved. We ran base and pr back-to-back within each workload, under one lock, so any drift over time hits both branches equally. Measuring three branches as three separate sessions would map any drift between sessions — thermal state, other users coming and going, neighbouring die activity — directly onto the branch axis.

2. Card selection. The chart says dev 3-7, which is five devices of mixed parity. We measured this effect directly: on an odd device whose neighbouring die was busy, Case1 split into two modes near 23 ms and 51 ms; on an exclusively held even device the second mode disappeared entirely, with the maximum at 30.9 ms. It also matches which rungs you flagged [BIMODAL] — bgemm and Mid4 unimodal, Mid16 / Case2 / Case1 bimodal. That ordering is by run duration, which is what neighbour interference looks like: the longer a workload occupies the card, the more likely it overlaps someone else's work.

3. A genuine machine difference — but the pattern argues against it. A machine or environment difference should shift both branches the same way. It does not. Your eager sits below the minimum of our #1618 distribution on all six rungs (7–34% faster), while your main matches our base within 1–6% on Mid4 / Mid16 / Mid64 and sits above the maximum of our base on Case2 and Case1. The two branches relate to our data by different rules, and a uniform machine offset cannot produce that.

Could you describe your setup in more detail?

Two notes on our side first, since both are places where we may already have diverged from you:

• Mid4 / Mid16 / Mid64 are our own reconstruction. They are not in Add: host_build_graph copies of the paged_attention examples #1667, so we derived them from the paged_attention case: batch scaled to 4 / 16 / 64 with Case1's other parameters left unchanged. The task counts come out at 1,028 / 4,112 / 16,448, which match your ladder — but if your rungs were built some other way, that is one place our numbers could differ from yours.
• qwen. Add: host_build_graph copy of the qwen3_14b_decode example #1666 adds the host_build_graph copy of qwen3_14b_decode, but we do not see it anywhere in your data, so we did not measure it this round either.

Beyond that: to work out where the difference comes from, could you describe your measurement in detail — the environment and how the runs were organised (branch builds, CANN version, which device(s) and whether they were held exclusively, and how the three branches were ordered relative to one another) — and the data collection point: which markers or which tool produce the "whole-device wall clock" and the "AICore task execution time" spans.

With that we can align our runs to yours and narrow down what produces the gap, rather than continuing to measure separately. Happy to re-run anything under a configuration you specify.